# -*- coding: utf-8 -*-

import requests
import json


API_KEY = "FBD3HDYzBNY0fHGFhpOOtR8I"
SECRET_KEY = "dFvdvjUUaEcXlBlgVhzYGLvc8yx5yTq1"

language_dict = {"en": "英语", "fr": "法语", "es": "西班牙语", "cn": "中文"}
sum_prompt_dict = {"enumeration": "你是一个专业会议记录者，总结我所说的话，要求：总结内容必须详细完整，输出结果用数字序号分点分段列出，不能输出与内容无关的句子。", 
                   "direct": "总结我的输入，不需要其他输出。", 
                   "concise": "你是一个专业会议记录者，用几个关键词组总结我所说的话，要求：内容用最少的词语高度概括我的话，不能输出与会议内容无关的句子。",
                  }

cn_trans_text = "谁的尾巴长？谁的尾巴短？谁的尾巴好像一把伞？猴子的尾巴长。兔子的尾巴短。松鼠的尾巴好像一把伞。谁的谁的尾巴弯？谁的尾扁？谁的尾巴最好看？公鸡的尾巴弯。鸭子的尾巴扁。孔雀的尾巴最好看。"
en_trans_text = "Who has a long tail? Who has a short tail? Who's tail is like an umbrella? Monkeys have a long tail. rabbits have a short tail. squirrels have a curly tail. Who has a curved tail? Who has a flat tail? Who's tail is the most beautiful? chickens have a curved tail. Ducks have a flat tail. Peacocks have the most beautiful tail."
long_text = "资料部祁主任：应该看到的是，办公软件虽然很多，但从专业角度而言，大都不很规范。我指的是编辑方面的问题、如WORD中对于行政公文这一块就干脆忽略掉，而书信这一部份也大多是英文习惯，中国人使用起来很不方便。WPS是中国人开发的软件，在技术上很有特点，但中国运用文方面的编辑十分简陋，离专业水准很远、我认为我们定位在这一方面是很有市场的。"
dialog_text = "技术部朱总：“类似的办公软件已经有不少，如微软公司的系统，金山公司的WPS系列，以及众多的财务，税务，管理方面的软件。我认为首要的问题是确定选题方向，如果没有特点，千万不能动手。”\n资料部祁主任：“应该看到的是，办公软件虽然很多，但从专业角度而言，大都不很规范。我指的是编辑方面的问题、如WORD中对于行政公文这一块就干脆忽略掉，而书信这一部份也大多是英文习惯，中国人使用起来很不方便。WPS是中国人开发的软件，在技术上很有特点，但中国运用文方面的编辑十分简陋，离专业水准很远、我认为我们定位在这一方面是很有市场的。”\n市场部唐主任：“这是在众多航空中间寻求突破，我认为有成功的希望，关键的问题就是必须小巧，并且速度极快、因为我们建造的不是航空，这就必须考虑到兼容问题。”"


def get_access_token():
    """
    使用 AK，SK 生成鉴权签名（Access Token）
    :return: access_token，或是None(如果错误)
    """
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    return str(requests.post(url, params=params).json().get("access_token"))

def get_ernie_response(main_prompt, dialog_prompt, text):
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + get_access_token()
    
    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": main_prompt
            },
            {
                "role": "assistant",
                "content": dialog_prompt
            },
            {
                "role": "user",
                "content": text
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    try:
        response = requests.request("POST", url, headers=headers, data=payload)
        response_text = eval(response.text.replace("false", "False"))["result"]
    except:
        response_text = "error"

    return response_text

def translation(text, dest_lan="en"):
    language = language_dict[dest_lan]
    main_prompt = f"你是一个{language}母语者，将句子全部翻译为{language}，要求：没有句子首字母大小写错误，翻译结果与原句意思必须完全对应，只输出{language}，不需要输出其他内容。"
    dialog_prompt = "好的，请问你需要翻译什么样的句子？"
    
    return get_ernie_response(main_prompt, dialog_prompt, text)

def summarization(text, style="direct"):
    main_prompt = sum_prompt_dict[style]
    dialog_prompt = "好的，请问你有什么需要总结的内容呢？"

    ernie_response = get_ernie_response(main_prompt, dialog_prompt, text)

    if style == "enumeration":
        index = ernie_response.index("1")
        ernie_response = ernie_response[index:]
    
    return ernie_response


if __name__ == '__main__':
    # 翻译
    # print(translation(cn_trans_text, "fr"))

    # 直接总结
    # print(summarization(long_text, "direct"))

    # 精炼总结
    # print(summarization(long_text, "concise"))

    # 分点总结
    # print(summarization(long_text, "enumeration"))

    # 对话总结 对话输入由前端获取，不同对话人用冒号和引号引出自己的说话内容，换行符分隔，合并为一个字符串
    print(summarization(dialog_text, "direct"))
    # print(summarization(dialog_text, "enumeration"))
